# for gcd function (or easily implementable to avoid import)
import fractions
# for random elements drawing in recover_prime_factors
import random
import pyasn1_modules.rfc3447
import recover
import gmpy
import pyasn1
import pyasn1.codec.ber.encoder
import base64


def fail_function():
	print("Prime factors not found")


def output_primes(a, n):
	p = fractions.gcd(a, n)
	q = int(n // p)
	if p > q:
		p, q = q, p
	print("Found factors p and q")
	print("p = {0}".format(str(p)))
	print("q = {0}".format(str(q)))
	return p, q


# 通过n e d 反推p, q
def recover_prime_factors(n, e, d):
	"""The following algorithm recovers the prime factor
		s of a modulus, given the public and private
		exponents.
		Function call: RecoverPrimeFactors(n, e, d)
		Input: 	n: modulus
				e: public exponent
				d: private exponent
		Output: (p, q): prime factors of modulus"""

	k = d * e - 1
	if k % 2 == 1:
		fail_function()
		return 0, 0
	else:
		t = 0
		r = k
		while r % 2 == 0:
			r = int(r // 2)
			t += 1
		for i in range(1, 101):
			# random g in [0, n-1]
			g = random.randint(0, n)
			y = pow(g, r, n)
			if y == 1 or y == n - 1:
				continue
			else:
				# j \in [1, t-1]
				for j in range(1, t):
					x = pow(y, 2, n)
					if x == 1:
						p, q = output_primes(y - 1, n)
						return p, q
					elif x == n - 1:
						continue
					y = x
					x = pow(y, 2, n)
					if x == 1:
						p, q = output_primes(y - 1, n)
						return p, q


# 生成Rsa私钥
def asn1_encode_priv_key(N, e, d, p, q):
      key = pyasn1_modules.rfc3447.RSAPrivateKey()
      dp = d % (p - 1)
      dq = d % (q - 1)
      qInv = gmpy.invert(q, p)
      #assert (qInv * q) % p == 1
      key.setComponentByName('version', 0)
      key.setComponentByName('modulus', N)
      key.setComponentByName('publicExponent', e)
      key.setComponentByName('privateExponent', d)
      key.setComponentByName('prime1', p)
      key.setComponentByName('prime2', q)
      key.setComponentByName('exponent1', dp)
      key.setComponentByName('exponent2', dq)
      key.setComponentByName('coefficient', qInv)
      ber_key = pyasn1.codec.ber.encoder.encode(key)
      pem_key = base64.b64encode(ber_key).decode("ascii")
      out = ['-----BEGIN RSA PRIVATE KEY-----']
      out += [pem_key[i:i + 64] for i in range(0, len(pem_key), 64)]
      out.append('-----END RSA PRIVATE KEY-----\n')
      out = "\n".join(out)
      return out.encode("ascii")


if __name__=='__main__':
	n = 0x3A6160848FB1734CBD0FA22CEF582E849223AC04510D51502556B6476D07397F03DF155289C20112E87C6F35361D9EB622CA4A0E52D9CD87BF723526C826B88387D06ABC4279E353F12AD8EC62EA73C47321A20B89644889A792A73152BC7014B80A693D2E58B123FA925C356B1EBA037A4DCAC8D8DE809167A6FCC30C5C785
	e = 0x0365962e8daba7ba92fc08768a5f73b3854f4c79969d5518a078a034437c4669bdb705be4d8b8babf4fda1a6e715269e87b28eecb0d4e02726a27fb8721863740720f583688e5567eb10729bb0d92b322d719949e40c57198d764f1c633e5e277da3d3281ece2ce2eb4df945be5afc3e78498ed0489b2459059664fe15c88a33
	d = 89508186630638564513494386415865407147609702392949250864642625401059935751367507

	p, q = recover.RecoverPrimeFactors(n, e, d)
	print asn1_encode_priv_key(n, e, d, p, q)